#!/usr/bin/python2
# -*- coding: utf-8 -*-
# Copyright (c) 2011 Thorsten Weimann <thorsten.weimann (at) gmx.net>
# Based on the Python AES implementation of Brandon Sterne
# http://brandon.sternefamily.net/files/pyAES.txt
# Licensed under the MIT license.

from __future__ import print_function

import os

from hashlib import sha256
try:
    from cStringIO import StrinIO
except ImportError:
    from StringIO import StringIO


#Rijndael S-Box
SBOX =  (0x63, 0x7c, 0x77, 0x7b, 0xf2, 0x6b, 0x6f, 0xc5, 0x30, 0x01, 0x67,
         0x2b, 0xfe, 0xd7, 0xab, 0x76, 0xca, 0x82, 0xc9, 0x7d, 0xfa, 0x59,
         0x47, 0xf0, 0xad, 0xd4, 0xa2, 0xaf, 0x9c, 0xa4, 0x72, 0xc0, 0xb7,
         0xfd, 0x93, 0x26, 0x36, 0x3f, 0xf7, 0xcc, 0x34, 0xa5, 0xe5, 0xf1,
         0x71, 0xd8, 0x31, 0x15, 0x04, 0xc7, 0x23, 0xc3, 0x18, 0x96, 0x05,
         0x9a, 0x07, 0x12, 0x80, 0xe2, 0xeb, 0x27, 0xb2, 0x75, 0x09, 0x83,
         0x2c, 0x1a, 0x1b, 0x6e, 0x5a, 0xa0, 0x52, 0x3b, 0xd6, 0xb3, 0x29,
         0xe3, 0x2f, 0x84, 0x53, 0xd1, 0x00, 0xed, 0x20, 0xfc, 0xb1, 0x5b,
         0x6a, 0xcb, 0xbe, 0x39, 0x4a, 0x4c, 0x58, 0xcf, 0xd0, 0xef, 0xaa,
         0xfb, 0x43, 0x4d, 0x33, 0x85, 0x45, 0xf9, 0x02, 0x7f, 0x50, 0x3c,
         0x9f, 0xa8, 0x51, 0xa3, 0x40, 0x8f, 0x92, 0x9d, 0x38, 0xf5, 0xbc,
         0xb6, 0xda, 0x21, 0x10, 0xff, 0xf3, 0xd2, 0xcd, 0x0c, 0x13, 0xec,
         0x5f, 0x97, 0x44, 0x17, 0xc4, 0xa7, 0x7e, 0x3d, 0x64, 0x5d, 0x19,
         0x73, 0x60, 0x81, 0x4f, 0xdc, 0x22, 0x2a, 0x90, 0x88, 0x46, 0xee,
         0xb8, 0x14, 0xde, 0x5e, 0x0b, 0xdb, 0xe0, 0x32, 0x3a, 0x0a, 0x49,
         0x06, 0x24, 0x5c, 0xc2, 0xd3, 0xac, 0x62, 0x91, 0x95, 0xe4, 0x79,
         0xe7, 0xc8, 0x37, 0x6d, 0x8d, 0xd5, 0x4e, 0xa9, 0x6c, 0x56, 0xf4,
         0xea, 0x65, 0x7a, 0xae, 0x08, 0xba, 0x78, 0x25, 0x2e, 0x1c, 0xa6,
         0xb4, 0xc6, 0xe8, 0xdd, 0x74, 0x1f, 0x4b, 0xbd, 0x8b, 0x8a, 0x70,
         0x3e, 0xb5, 0x66, 0x48, 0x03, 0xf6, 0x0e, 0x61, 0x35, 0x57, 0xb9,
         0x86, 0xc1, 0x1d, 0x9e, 0xe1, 0xf8, 0x98, 0x11, 0x69, 0xd9, 0x8e,
         0x94, 0x9b, 0x1e, 0x87, 0xe9, 0xce, 0x55, 0x28, 0xdf, 0x8c, 0xa1,
         0x89, 0x0d, 0xbf, 0xe6, 0x42, 0x68, 0x41, 0x99, 0x2d, 0x0f, 0xb0,
         0x54, 0xbb, 0x16)

# Rijndael Inverted S-Box
ISBOX = (0x52, 0x09, 0x6a, 0xd5, 0x30, 0x36, 0xa5, 0x38, 0xbf, 0x40, 0xa3,
         0x9e, 0x81, 0xf3, 0xd7, 0xfb , 0x7c, 0xe3, 0x39, 0x82, 0x9b, 0x2f,
         0xff, 0x87, 0x34, 0x8e, 0x43, 0x44, 0xc4, 0xde, 0xe9, 0xcb , 0x54,
         0x7b, 0x94, 0x32, 0xa6, 0xc2, 0x23, 0x3d, 0xee, 0x4c, 0x95, 0x0b,
         0x42, 0xfa, 0xc3, 0x4e , 0x08, 0x2e, 0xa1, 0x66, 0x28, 0xd9, 0x24,
         0xb2, 0x76, 0x5b, 0xa2, 0x49, 0x6d, 0x8b, 0xd1, 0x25 , 0x72, 0xf8,
         0xf6, 0x64, 0x86, 0x68, 0x98, 0x16, 0xd4, 0xa4, 0x5c, 0xcc, 0x5d,
         0x65, 0xb6, 0x92 , 0x6c, 0x70, 0x48, 0x50, 0xfd, 0xed, 0xb9, 0xda,
         0x5e, 0x15, 0x46, 0x57, 0xa7, 0x8d, 0x9d, 0x84 , 0x90, 0xd8, 0xab,
         0x00, 0x8c, 0xbc, 0xd3, 0x0a, 0xf7, 0xe4, 0x58, 0x05, 0xb8, 0xb3,
         0x45, 0x06 , 0xd0, 0x2c, 0x1e, 0x8f, 0xca, 0x3f, 0x0f, 0x02, 0xc1,
         0xaf, 0xbd, 0x03, 0x01, 0x13, 0x8a, 0x6b , 0x3a, 0x91, 0x11, 0x41,
         0x4f, 0x67, 0xdc, 0xea, 0x97, 0xf2, 0xcf, 0xce, 0xf0, 0xb4, 0xe6,
         0x73 , 0x96, 0xac, 0x74, 0x22, 0xe7, 0xad, 0x35, 0x85, 0xe2, 0xf9,
         0x37, 0xe8, 0x1c, 0x75, 0xdf, 0x6e , 0x47, 0xf1, 0x1a, 0x71, 0x1d,
         0x29, 0xc5, 0x89, 0x6f, 0xb7, 0x62, 0x0e, 0xaa, 0x18, 0xbe, 0x1b ,
         0xfc, 0x56, 0x3e, 0x4b, 0xc6, 0xd2, 0x79, 0x20, 0x9a, 0xdb, 0xc0,
         0xfe, 0x78, 0xcd, 0x5a, 0xf4 , 0x1f, 0xdd, 0xa8, 0x33, 0x88, 0x07,
         0xc7, 0x31, 0xb1, 0x12, 0x10, 0x59, 0x27, 0x80, 0xec, 0x5f , 0x60,
         0x51, 0x7f, 0xa9, 0x19, 0xb5, 0x4a, 0x0d, 0x2d, 0xe5, 0x7a, 0x9f,
         0x93, 0xc9, 0x9c, 0xef , 0xa0, 0xe0, 0x3b, 0x4d, 0xae, 0x2a, 0xf5,
         0xb0, 0xc8, 0xeb, 0xbb, 0x3c, 0x83, 0x53, 0x99, 0x61 , 0x17, 0x2b,
         0x04, 0x7e, 0xba, 0x77, 0xd6, 0x26, 0xe1, 0x69, 0x14, 0x63, 0x55,
         0x21, 0x0c, 0x7d)

# Rijndael Rcon
RCON = (0x8d, 0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, 0x1b, 0x36,
        0x6c, 0xd8, 0xab, 0x4d, 0x9a, 0x2f, 0x5e, 0xbc, 0x63, 0xc6, 0x97,
        0x35, 0x6a, 0xd4, 0xb3, 0x7d, 0xfa, 0xef, 0xc5, 0x91, 0x39, 0x72,
        0xe4, 0xd3, 0xbd, 0x61, 0xc2, 0x9f, 0x25, 0x4a, 0x94, 0x33, 0x66,
        0xcc, 0x83, 0x1d, 0x3a, 0x74, 0xe8, 0xcb, 0x8d, 0x01, 0x02, 0x04,
        0x08, 0x10, 0x20, 0x40, 0x80, 0x1b, 0x36, 0x6c, 0xd8, 0xab, 0x4d,
        0x9a, 0x2f, 0x5e, 0xbc, 0x63, 0xc6, 0x97, 0x35, 0x6a, 0xd4, 0xb3,
        0x7d, 0xfa, 0xef, 0xc5, 0x91, 0x39, 0x72, 0xe4, 0xd3, 0xbd, 0x61,
        0xc2, 0x9f, 0x25, 0x4a, 0x94, 0x33, 0x66, 0xcc, 0x83, 0x1d, 0x3a,
        0x74, 0xe8, 0xcb, 0x8d, 0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40,
        0x80, 0x1b, 0x36, 0x6c, 0xd8, 0xab, 0x4d, 0x9a, 0x2f, 0x5e, 0xbc,
        0x63, 0xc6, 0x97, 0x35, 0x6a, 0xd4, 0xb3, 0x7d, 0xfa, 0xef, 0xc5,
        0x91, 0x39, 0x72, 0xe4, 0xd3, 0xbd, 0x61, 0xc2, 0x9f, 0x25, 0x4a,
        0x94, 0x33, 0x66, 0xcc, 0x83, 0x1d, 0x3a, 0x74, 0xe8, 0xcb, 0x8d,
        0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, 0x1b, 0x36, 0x6c,
        0xd8, 0xab, 0x4d, 0x9a, 0x2f, 0x5e, 0xbc, 0x63, 0xc6, 0x97, 0x35,
        0x6a, 0xd4, 0xb3, 0x7d, 0xfa, 0xef, 0xc5, 0x91, 0x39, 0x72, 0xe4,
        0xd3, 0xbd, 0x61, 0xc2, 0x9f, 0x25, 0x4a, 0x94, 0x33, 0x66, 0xcc,
        0x83, 0x1d, 0x3a, 0x74, 0xe8, 0xcb, 0x8d, 0x01, 0x02, 0x04, 0x08,
        0x10, 0x20, 0x40, 0x80, 0x1b, 0x36, 0x6c, 0xd8, 0xab, 0x4d, 0x9a,
        0x2f, 0x5e, 0xbc, 0x63, 0xc6, 0x97, 0x35, 0x6a, 0xd4, 0xb3, 0x7d,
        0xfa, 0xef, 0xc5, 0x91, 0x39, 0x72, 0xe4, 0xd3, 0xbd, 0x61, 0xc2,
        0x9f, 0x25, 0x4a, 0x94, 0x33, 0x66, 0xcc, 0x83, 0x1d, 0x3a, 0x74,
        0xe8, 0xcb)


def _rotate(word, n=1):
    """Returns a copy of the word shifted n bytes."""
    return word[n:] + word[:n]


def _shift_rows(state, inverse=False):
    """Iterates over each row and shift the bytes to the LEFT or to
    the RIGHT (if inverse = True) by the appropriate offset."""
    tmp = bytearray(len(state))
    for row, offset in enumerate(range(0, 16, 4)):
        if inverse:
            row = -row
        tmp[offset:offset+4] = _rotate(state[offset:offset+4], row)
    return tmp


def _schedule_core(word, iteration):
    word = _rotate(word, 1)
    new_word = bytearray()
    for byte in word:
        new_word.append(SBOX[byte])
    new_word[0] = new_word[0] ^ RCON[iteration]
    return new_word


def _expand_key(cipher_key):
    """Expands 256 bit cipher into 240 byte key."""
    key_size = len(cipher_key)
    expanded_key = cipher_key[:32]
    current_size = key_size
    rcon_iteration = 1
    tmp = bytearray(4)
    while current_size < 240:
        for i in range(4):
            tmp[i] = expanded_key[current_size-4+i]
        if current_size % key_size == 0:
            tmp = _schedule_core(tmp, rcon_iteration)
            rcon_iteration += 1
        elif current_size % key_size == 16:
            tmp = bytearray([SBOX[x] for x in tmp])
        for byte in tmp:
            expanded_key.append(expanded_key[current_size - key_size] ^ byte)
            current_size += 1
    return expanded_key


def _substitute_bytes(state, inverse=False):
    """Do sbox or inverse sbox transform on each byte on the state table."""
    box = SBOX if not inverse else ISBOX
    return bytearray([box[x] for x in state])


def _add_round_key(state, round_key):
    """XOR each byte of the round key with the state table."""
    return bytearray([x ^ y for x, y in zip(state, round_key)])


def galois_mul(a, b):
    """Galois multiplication of a and b."""
    p = 0
    hi_bit_set = 0
    for i in range(8):
        if b & 1 == 1:
            p ^= a
        hi_bit_set = a & 0x80
        a <<= 1
        if hi_bit_set == 0x80:
            a ^= 0x1b
        b >>= 1
    return p % 256


def _mix_column(column, inverse=False):
    f = (2, 1, 1, 3) if not inverse else (14, 9, 13, 11)
    col = bytearray(4)
    c = column
    g = galois_mul
    col[0] = g(c[0], f[0]) ^ g(c[3], f[1]) ^ g(c[2], f[2]) ^ g(c[1], f[3])
    col[1] = g(c[1], f[0]) ^ g(c[0], f[1]) ^ g(c[3], f[2]) ^ g(c[2], f[3])
    col[2] = g(c[2], f[0]) ^ g(c[1], f[1]) ^ g(c[0], f[2]) ^ g(c[3], f[3])
    col[3] = g(c[3], f[0]) ^ g(c[2], f[1]) ^ g(c[1], f[2]) ^ g(c[0], f[3])
    return col


def _mix_columns(state, inverse=False):
    """Wrapper for _mix_column()."""
    tmp = bytearray(len(state))
    for i in range(4):
        column = bytearray()
        for j in range(4):
            column.append(state[j*4+i])
        column = _mix_column(column, inverse)
        for j in range(4):
            tmp[j*4+i] = column[j]
    return tmp


def _aes_round(state, round_key):
    """Applies all four transformations in order."""
    tmp = state[:]
    tmp = _substitute_bytes(tmp)
    tmp = _shift_rows(tmp)
    tmp = _mix_columns(tmp)
    tmp = _add_round_key(tmp, round_key)
    return tmp


def _aes_round_inverse(state, round_key):
    """Applies all four inverse transformations."""
    tmp = state[:]
    tmp = _add_round_key(tmp, round_key)
    tmp = _mix_columns(tmp, True)
    tmp = _shift_rows(tmp, True)
    tmp = _substitute_bytes(tmp, True)
    return tmp


def _create_round_key(expanded_key, n):
    """Creates a 16-byte round key based on an expanded key and
    round number."""
    return expanded_key[n*16:n*16+16]


def _password_to_key(password):
    """Creates a key from a user-supplied password."""
    if callable(password):
        password = password().strip()
    return bytearray(sha256(password).digest())


def _aes_main(state, expanded_key, rounds=14):
    tmp = state[:]
    round_key = _create_round_key(expanded_key, 0)
    tmp = _add_round_key(tmp, round_key)
    for i in range(1, rounds):
        round_key = _create_round_key(expanded_key, i)
        tmp = _aes_round(tmp, round_key)
    round_key = _create_round_key(expanded_key, rounds)
    tmp = _substitute_bytes(tmp)
    tmp = _shift_rows(tmp)
    tmp = _add_round_key(tmp, round_key)
    return tmp


def _aes_main_inverse(state, expanded_key, rounds=14):
    tmp = state[:]
    round_key = _create_round_key(expanded_key, rounds)
    tmp = _add_round_key(tmp, round_key)
    tmp = _shift_rows(tmp, True)
    tmp = _substitute_bytes(tmp, True)
    for i in range(rounds-1, 0, -1):
        round_key = _create_round_key(expanded_key, i)
        tmp = _aes_round_inverse(tmp, round_key)
    round_key = _create_round_key(expanded_key, 0)
    tmp = _add_round_key(tmp, round_key)
    return tmp


def _aes_encrypt(plaintext, key):
    """Encrypt a single block of plaintext."""
    return _aes_main(plaintext, _expand_key(key))


def _aes_decrypt(ciphertext, key):
    """Decrypt a single block of ciphertext."""
    return _aes_main_inverse(ciphertext, _expand_key(key))


def _get_block(fp):
    """Returns 16-byte block from an open file or file like object."""
    bytes = fp.read(16)
    if not bytes:
        return ''
    block = bytearray([ord(x) for x in bytes])
    if len(block) < 16:
        pad_char = 16 - len(block)
        while len(block) < 16:
            block.append(pad_char)
    return block


def _encrypt(fp, password, out):
    """Encrypts by reading from `fp` and writing to `out`."""
    iv = bytearray(os.urandom(16))
    aes_key = _password_to_key(password)
    for byte in iv:
        out.write(chr(byte))
    fp.seek(0, 2)
    size = fp.tell()
    fp.seek(0)
    first_round = True
    block = _get_block(fp)
    while block:
        if first_round:
            block_key = _aes_encrypt(iv, aes_key)
            first_round = False
        else:
            block_key = _aes_encrypt(block_key, aes_key)
        ciphertext = bytearray([x ^ y for x, y in zip(block, block_key)])
        for byte in ciphertext:
            out.write(chr(byte))
        block = _get_block(fp)
    if size % 16 == 0:
        out.write(16*chr(16))
    fp.close()


def _decrypt(fp, password, out):
    """Decrypts by reading from `fp` and writing to `out`."""
    aes_key = _password_to_key(password)
    iv = _get_block(fp)
    fp.seek(0, 2)
    size = fp.tell()
    fp.seek(16)
    first_round = True
    block = _get_block(fp)
    while block:
        if first_round:
            block_key = _aes_encrypt(iv, aes_key)
            first_round = False
        else:
            block_key = _aes_encrypt(block_key, aes_key)
        plaintext = bytearray([x ^ y for x, y in zip(block, block_key)])
        if fp.tell() == size:
            plaintext = plaintext[:-plaintext[-1]]
        for byte in plaintext:
            out.write(chr(byte))
        block = _get_block(fp)
    fp.close()


def encrypt_text(text, password):
    """Encrypts the given text using password."""
    fp = StringIO(text)
    out = StringIO()
    _encrypt(fp, password, out)
    try:
        return out.getvalue()
    finally:
        out.close()


def encrypt_file(filename, password, outfile=None):
    if outfile is None:
        outfile = '{0}.aes'.format(filename)
    fp = open(filename, 'rb')
    with open(outfile, 'wb') as out:
        _encrypt(fp, password, out)


def decrypt_text(ciphertext, password):
    fp = StringIO(ciphertext)
    out = StringIO()
    _decrypt(fp, password, out)
    try:
        return out.getvalue()
    finally:
        out.close()


def decrypt_file(filename, password, outfile=None):
    if outfile is None:
        outfile = '{0}.plain'.format(filename)
    fp = open(filename, 'rb')
    with open(outfile, 'wb') as out:
        _decrypt(fp, password, out)


encrypt = encrypt_text
decrypt = decrypt_text


# Some tests
if __name__ == '__main__':
    import unittest

    class AESTest(unittest.TestCase):
        def test_string(self):
            s = 'This is a little test.'
            encrypted = encrypt(s, 'secret')
            self.assertEqual(s, decrypt(encrypted, 'secret'))

        def test_utf8(self):
            u = u'This is a test with german umlauts: äöü'.encode('utf-8')
            encrypted = encrypt(u, 'secret')
            self.assertEqual(u, decrypt(encrypted, 'secret'))

        def test_passwords(self):
            encrypted1 = encrypt('Test', 'secret')
            encrypted2 = encrypt('Test', 'Secret')
            self.assertNotEqual(encrypted1, encrypted2)

        def test_content(self):
            encrypted1 = encrypt('test', 'secret')
            encrypted2 = encrypt('Test', 'secret')
            self.assertNotEqual(encrypted1, encrypted2)

    suite = unittest.TestLoader().loadTestsFromTestCase(AESTest)
    unittest.TextTestRunner(verbosity=2).run(suite)


